from solvers.solverbase import SolverBase
from solvers.docarmo.docarmobase import DoCarmoBase
from solvers.stabilizations import supg
from dolfin import *
import numpy

class CrankNicolson(SolverBase,DoCarmoBase):
    'Crank-Nicolson method with nonlinear-supg stabilization.'
    def __init__(self,p=1):
        SolverBase.__init__(self,p)
        DoCarmoBase.__init__(self)

    def solve(self,problem):
        mesh = problem.domain

        V = FunctionSpace(mesh,'CG',1)
        VV = VectorFunctionSpace(mesh,'CG',1)
        DGV = FunctionSpace(mesh,'DG',0)
        DGVV = VectorFunctionSpace(mesh,'DG',0)
  
        #-----------------------------------------------------------------------------------------------------------------------------
        def projection(v,phi,phi_,p_,dt,cache): # velocity,current solution,previous solution,projection(output)
            p = TrialFunction(DGVV)
            phiT = TestFunction(DGVV)

            if 'aVV' not in cache.keys():
                a = inner(p,phiT)*dx
                cache['aVV'] = assemble(a)

            w = conditional(
                        gt(sqrt(inner(nabla_grad(phi),nabla_grad(phi))),1E-16),\
                        (phi-phi_+dt*inner(v,nabla_grad(phi)))*nabla_grad(phi)/inner(nabla_grad(phi),nabla_grad(phi)),\
                        Constant((0.,0.))\
                           ) 

            l = inner(w,phiT)*dx
            b = assemble(l)
            solve(cache['aVV'],p_.vector(),b)
            
        def get_sigma(v,u,u_,p_,dt,sigma_,cache): # velocity,current sol,previous level solution,projection,coefs(output) 
            trial = TrialFunction(DGV)
            psi = TestFunction(DGV)

            ctrial = TrialFunction(V)
            cpsi = TestFunction(V)
            
            norm_v = Function(V)
            norm_p = Function(DGV)
            resF = Function(DGV)
            
            if 'aV' not in cache.keys():
                a =  trial*psi*dx
                cache['aV'] = assemble(a)
            
            if 'caV' not in cache.keys():
                ca = ctrial*cpsi*dx
                cache['caV'] = assemble(ca)

            #-----
            p_norm_l = sqrt(inner(p_,p_))*psi*dx
            b1 = assemble(p_norm_l)
            solve(cache['aV'],norm_p.vector(),b1)               # get the norm of p_(discon)
            #-----
            v_norm_l = sqrt(inner(v,v))*cpsi*dx
            b2 = assemble(v_norm_l)
            solve(cache['caV'],norm_v.vector(),b2)              # get the norm of v (continuous) ! 
            #-----
            resl = (u-u_+dt*inner(v,nabla_grad(u)))*psi*dx
            b3 = assemble(resl)
            solve(cache['aV'],resF.vector(),b3)                 # get the residuum
            
            for i in xrange(len(sigma_.vector().array())):
                c = Cell(mesh,i)
                x = c.midpoint()

                v_mag = norm_v(x.x(),x.y())
                p_mag = norm_p(x.x(),x.y())
                residuum = resF(x.x(),x.y())
                
                if abs(residuum) < 1E-16:
                    sigma_.vector()[i] = 0
                else:
                    alfa = 1.
                    if p_mag >= v_mag:
                        sigma_.vector()[i] = 0
                    else:
                        sigma_.vector()[i] = v_mag*p_mag - alfa*p_mag**2 
        #----------------------------------------------------------------------------------------------------------------------------------
       
        phi_ = interpolate(problem.phi_,V)
        phi = TrialFunction(V)
        varphi = TestFunction(V)

        t = problem.t           
        dt = Constant(SolverBase._get_time_step(self,problem,VV,t))
        T = problem.T

        v = SolverBase._get_velocity(self,problem,VV,t)
        tau = Function(DGV)
        tau.vector()[:] = supg.get_tau(v,eps=supg.diffusivity)

        sigma = Function(DGV)
        sigma_ = Function(DGV)

        self.update(problem,phi_,t,float(dt),0)  # store the solution at time 0
        
        pert = conditional(
                           gt(sqrt(inner(v,v)),1E-16),
                           Constant(1.0),
                           Constant(0.) 
                           ) # a way to realize swith on |v| == 0.
                        
        a = phi*varphi*dx + 0.5*dt*inner(v,nabla_grad(phi))*varphi*dx +\
            tau*pert*phi*inner(v,nabla_grad(varphi))*dx +\
            0.5*dt*tau*pert*inner(dot(nabla_grad(phi),outer(v,v)),nabla_grad(varphi))*dx+\
            0.5*dt*tau*sigma*inner(nabla_grad(phi),nabla_grad(varphi))*dx

        l = phi_*varphi*dx + tau*pert*phi_*inner(v,nabla_grad(varphi))*dx\
            - 0.5*dt*inner(v,nabla_grad(phi_))*varphi*dx+\
            -0.5*dt*tau*pert*inner(dot(nabla_grad(phi_),outer(v,v)),nabla_grad(varphi))*dx\
            -0.5*dt*tau*sigma*inner(nabla_grad(phi_),nabla_grad(varphi))*dx
        
        t = 0
        count = 0
        maxIter = 10
        tol = 1E-3
        eps = 1E-8                               # viscosity TODO is it needed? what values are used in literature

        cache = {}
        phi = Function(V)           # solution
        phi_m = Function(V)         # solution in iteration loop
        p_ = Function(DGVV)         # projection
        while t <= T:
            t += float(dt)
            count += 1
            iter = 0
            error = 1

            projection(v,phi_,phi_,p_,dt,cache)       # guess
            get_sigma(v,phi_,phi_,p_,dt,sigma,cache)  # get the sigma for first iteration with phi_0 = phi_(previous solution)
            get_sigma(v,phi_,phi_,p_,dt,sigma_,cache) # for the RHS
            L = assemble(l) 
            while error > tol and iter < maxIter:
                iter += 1
                A = assemble(a)     
                solve(A,phi.vector(),L)
                diff = phi.vector().array() - phi_m.vector().array()
                error = numpy.linalg.norm(diff,ord=numpy.Inf)
                phi_m.assign(phi)
                projection(v,phi_m,phi_,p_,dt,cache)       # proper
                get_sigma(v,0.5*(phi_m+phi_),phi_,p_,dt,sigma,cache)  # get the sigma for the next iteration with phi_j = phi(current solution)
                #print iter, error

            phi_.assign(phi)
            self.update(problem,phi_,t,float(dt),count)

        problem.save_data(self.saveDir,str(self.CFL))
